use std::env;
use std::io::{Error, ErrorKind, Write};
use std::iter::once;
use std::path::Path;
use std::process::{Command, Stdio};
use std::ptr;
use winapi::shared::minwindef::FALSE;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::handleapi::CloseHandle;
use winapi::um::processthreadsapi::OpenProcessToken;
use winapi::um::processthreadsapi::{CreateProcessAsUserW, GetCurrentProcess};
use winapi::um::processthreadsapi::{PROCESS_INFORMATION, STARTUPINFOA};
use winapi::um::securitybaseapi::CreateRestrictedToken;
use winapi::um::winbase::{CREATE_UNICODE_ENVIRONMENT, STARTF_USESHOWWINDOW};
use winapi::um::winnt::{HANDLE, SANDBOX_INERT, TOKEN_ADJUST_DEFAULT};
use winapi::um::winnt::{TOKEN_ASSIGN_PRIMARY, TOKEN_DUPLICATE, TOKEN_QUERY};
use winapi::um::winuser::SW_SHOWNORMAL;

/// 创建降级权限的进程令牌, 返回受限令牌句柄.
/// 因为PostgreSQL禁止使用管理员权限，一般Windows用户会带有管理员身份，
/// 因此需要自动对启动进程进行权限降级。The server must be started  under
/// an unprivileged user ID to prevent possible system security compromises.
fn create_restricted_token() -> std::io::Result<HANDLE> {
    // 获取当前进程令牌
    let mut token_handle = ptr::null_mut();
    let result = unsafe {
        OpenProcessToken(
            GetCurrentProcess(),
            TOKEN_DUPLICATE | TOKEN_ADJUST_DEFAULT | TOKEN_QUERY | TOKEN_ASSIGN_PRIMARY,
            &mut token_handle,
        )
    };

    if result == 0 {
        return Err(Error::last_os_error());
    }

    // 创建受限令牌
    let mut restricted_token_handle = ptr::null_mut();
    let result = unsafe {
        CreateRestrictedToken(
            token_handle,
            SANDBOX_INERT,
            0,
            ptr::null_mut(),
            0,
            ptr::null_mut(),
            0,
            ptr::null_mut(),
            &mut restricted_token_handle,
        )
    };

    unsafe {
        CloseHandle(token_handle);
    }

    if result == 0 {
        return Err(Error::last_os_error());
    }

    Ok(restricted_token_handle)
}

// 以降级权限启动PostgreSQL
fn start_postgres(current_dir: &Path) -> std::io::Result<()> {
    // PostgreSQL可执行文件路径
    let bin_dir = current_dir.join("bin");
    let postgres_exe = bin_dir.join("postgres.exe");
    let data_dir = current_dir.join("pgdata");

    let postgres_path = postgres_exe.to_string_lossy().to_string();
    let data_dir_path = data_dir.to_string_lossy().to_string();

    // 确保路径存在
    if !postgres_exe.exists() || !postgres_exe.is_file() {
        println!("错误: 找不到PostgreSQL可执行文件: {}", postgres_path);
        return Err(Error::new(
            ErrorKind::NotFound,
            "找不到PostgreSQL可执行文件",
        ));
    }

    // 准备启动信息和进程信息
    let mut startup_info: STARTUPINFOA = unsafe { std::mem::zeroed() };
    startup_info.cb = std::mem::size_of::<STARTUPINFOA>() as u32;
    startup_info.dwFlags = STARTF_USESHOWWINDOW;
    startup_info.wShowWindow = SW_SHOWNORMAL as u16;

    let mut process_info: PROCESS_INFORMATION = unsafe { std::mem::zeroed() };

    let cmd_line = vec![
        postgres_path.as_str(),
        "-D",
        data_dir_path.as_str(),
        "-d",
        "1",
        "-E",
        "-F",
        "-c",
        "log_line_prefix=",
    ]
    .into_iter()
    .map(|s| {
        if s.contains(' ') {
            format!("\"{}\"", s) // Quote paths with spaces
        } else {
            s.to_string()
        }
    })
    .collect::<Vec<String>>()
    .join(" ");

    // 使用W版本的API，接受宽字符串命令行，需要转换命令行为宽字符
    let wide_cmd: Vec<u16> = cmd_line.encode_utf16().chain(once(0)).collect();

    // 获取受限令牌
    let restricted_token = create_restricted_token()?;

    // 启动进程
    let result = unsafe {
        CreateProcessAsUserW(
            restricted_token,
            ptr::null(),
            wide_cmd.as_ptr() as *mut _,
            ptr::null_mut(),
            ptr::null_mut(),
            FALSE,
            CREATE_UNICODE_ENVIRONMENT,
            ptr::null_mut(),
            ptr::null(),
            &mut startup_info as *mut STARTUPINFOA as *mut _,
            &mut process_info,
        )
    };

    // 关闭令牌句柄
    unsafe {
        CloseHandle(restricted_token);
    }

    if result == 0 {
        let error_code = unsafe { GetLastError() };
        println!("启动PostgreSQL失败: 错误代码 {}", error_code);
        return Err(Error::from_raw_os_error(error_code as i32));
    }

    // 关闭进程和线程句柄
    unsafe {
        CloseHandle(process_info.hProcess);
        CloseHandle(process_info.hThread);
    }

    Ok(())
}

fn run() -> std::io::Result<()> {
    let home_dir = match env::current_exe()?.parent() {
        Some(dir) => dir.to_path_buf(),
        None => {
            return Err(Error::new(ErrorKind::Other, "无法获取可执行文件目录"));
        }
    };

    println!("PostgreSQL Root: {}", home_dir.to_string_lossy());

    // 检查数据目录是否存在
    let data_dir = home_dir.join("pgdata");
    if !data_dir.exists() {
        initialize_pgdata(&home_dir)?;
    }

    if !data_dir.is_dir() {
        return Err(Error::new(
            ErrorKind::Other,
            format!("pgdata不是一个目录: {}", data_dir.to_string_lossy()),
        ));
    }

    let pg_version_file = data_dir.join("PG_VERSION");
    if !pg_version_file.exists() {
        return Err(Error::new(
            ErrorKind::Other,
            format!("无效数据空间: {}", data_dir.to_string_lossy()),
        ));
    }

    // 检查bin目录是否存在
    let bin_dir = home_dir.join("bin");
    if !bin_dir.exists() || !bin_dir.is_dir() {
        return Err(Error::new(
            ErrorKind::NotFound,
            format!("找不到PostgreSQL bin目录: {}", bin_dir.to_string_lossy()),
        ));
    }

    // 检查postgres.exe是否存在
    let postgres_exe = bin_dir.join("postgres.exe");
    if !postgres_exe.exists() || !postgres_exe.is_file() {
        return Err(Error::new(
            ErrorKind::NotFound,
            format!("找不到postgres.exe: {}", postgres_exe.to_string_lossy()),
        ));
    }

    // 启动PostgreSQL数据库
    start_postgres(&home_dir)?;

    Ok(())
}

fn main() {
    std::io::stdout().flush().unwrap();
    if let Err(e) = run() {
        eprintln!("\n错误: {:?}", e);
        eprintln!("按Enter键退出...");
        let _ = std::io::stdout().flush();
        let mut buffer = String::new();
        let _ = std::io::stdin().read_line(&mut buffer);
    }
}

fn initialize_pgdata(home_dir: &Path) -> std::io::Result<()> {
    let bin_dir = home_dir.join("bin");
    let initdb_path = bin_dir.join("initdb.exe");

    println!("Initializing PostgreSQL data space...");

    let initdb_status = Command::new(&initdb_path)
        .arg("-D")
        .arg("pgdata")
        .arg("-U")
        .arg("postgres")
        .arg("-A")
        .arg("trust")
        .arg("--encoding")
        .arg("UTF8")
        .arg("--locale")
        .arg("zh_CN")
        .arg("--no-instructions")
        .env("PGCLIENTENCODING", "GBK")
        .current_dir(home_dir)
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .status()
        .map_err(|e| {
            Error::new(
                ErrorKind::Other,
                format!("Failed to run initdb at {}: {}", initdb_path.display(), e),
            )
        })?;

    if !initdb_status.success() {
        return Err(Error::new(
            ErrorKind::Other,
            "Failed to initialize PostgreSQL with initdb",
        ));
    }
    println!("PostgreSQL initialized successfully.\n");

    Ok(())
}
